import cv2
import numpy as np
import onnxruntime as ort

# 配置参数
INPUT_SIZE = 640  # 与导出ONNX时的尺寸一致
# -------------------------- 类别定义 --------------------------
CLASSES = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
    "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
    "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
    "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
    "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
    "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
    "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
    "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote",
    "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book",
    "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
]  # 80 classe

# -------------------------- 1. 加载ONNX模型 --------------------------
session = ort.InferenceSession("yolov5s.onnx", providers=["CUDAExecutionProvider"])  # 使用GPU加速
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name


# -------------------------- 2. 预处理函数 --------------------------
def preprocess(image):
    # Letterbox 缩放 (保持宽高比)
    h, w = image.shape[:2]
    scale = min(INPUT_SIZE / h, INPUT_SIZE / w)
    nh, nw = int(h * scale), int(w * scale)
    resized = cv2.resize(image, (nw, nh))

    # 创建画布并填充边缘
    canvas = np.full((INPUT_SIZE, INPUT_SIZE, 3), 114, dtype=np.uint8)
    dh, dw = (INPUT_SIZE - nh) // 2, (INPUT_SIZE - nw) // 2
    canvas[dh:dh + nh, dw:dw + nw] = resized

    # 归一化并转置为 CHW 格式
    blob = canvas.astype(np.float32) / 255.0
    blob = blob.transpose(2, 0, 1)[np.newaxis, ...]  # [1,3,640,640]
    return blob, (scale, dw, dh)


# -------------------------- 3. 后处理函数 --------------------------
def postprocess(outputs, meta, conf_thres=0.5, iou_thres=0.5):
    scale, dw, dh = meta
    detections = outputs[0]  # 输出形状 [1,25200,85]

    # 过滤低置信度检测框
    mask = detections[..., 4] > conf_thres
    detections = detections[mask]

    # 转换坐标 (xywh to xyxy)
    boxes = detections[..., :4].copy()
    boxes[..., 0] = (boxes[..., 0] - dw) / scale  # x
    boxes[..., 1] = (boxes[..., 1] - dh) / scale  # y
    boxes[..., 2] = boxes[..., 2] / scale  # w
    boxes[..., 3] = boxes[..., 3] / scale  # h
    x1 = boxes[..., 0] - boxes[..., 2] / 2
    y1 = boxes[..., 1] - boxes[..., 3] / 2
    x2 = boxes[..., 0] + boxes[..., 2] / 2
    y2 = boxes[..., 1] + boxes[..., 3] / 2
    boxes = np.stack([x1, y1, x2, y2], axis=1)

    # 非极大值抑制 (NMS)
    scores = detections[..., 4]
    class_ids = np.argmax(detections[..., 5:], axis=-1)
    indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), conf_thres, iou_thres)

    # 添加类别ID校验
    valid_indices = [i for i in indices if class_ids[i] < len(CLASSES)]
    return boxes[valid_indices], scores[valid_indices], class_ids[valid_indices]


# -------------------------- 4. 推理流程 --------------------------
def infer_onnx(image_path, output_image_path):
    # 读取图像
    image = cv2.imread(image_path)

    # 预处理
    blob, meta = preprocess(image)

    # 执行推理
    outputs = session.run([output_name], {input_name: blob})

    # 后处理
    boxes, scores, class_ids = postprocess(outputs, meta)

    # 可视化结果
    for box, score, cls_id in zip(boxes, scores, class_ids):
        if cls_id >= len(CLASSES):
            continue  # 跳过无效类别
        x1, y1, x2, y2 = map(int, box)
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        label = f"{CLASSES[cls_id]}: {score:.2f}"
        cv2.putText(image, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

    # 保存处理后的图像
    cv2.imwrite(output_image_path, image)


if __name__ == "__main__":
    # 假设输入图像名为 "test.jpg" 并希望将结果保存为 "output.jpg"
    infer_onnx("test.jpg", "output_onnxruntime.jpg")